#%%
import pandas as pd, numpy as np, json, scipy.optimize as op, matplotlib.pyplot as plt, pathlib, os, lmfit as lm
%matplotlib inline

# Bucket calibration 
# Taylors calibration
# Bucket calibration is used to measure water flux from flume
h = np.array([np.array([17.9,  17.8, 17.8, 17.85]).mean(),
              np.array([17.7,  17.6, 17.8, 17.6]).mean(),
              np.array([19.75, 19.8, 19.8, 19.8]).mean(),
              np.array([19.9,  19.9, 19.9, 19.85]).mean(),
              np.array([21.0,  21.1, 21.0, 21.0]).mean(),
              np.array([21.8,  21.9, 21.8, 21.8]).mean()]).astype(float) # cm
vol = np.array([10.0, 10.0, 11.25, 11.25, 12.0, 12.5]) # L
# plt.plot(h, vol, 'ko')

def depth_vol(x, m, b):
    return m*x + b

(bucket_m, bucket_b), non = op.curve_fit(depth_vol, h, vol)
x = np.linspace(15, 25, 100)
# plt.plot(x, depth_vol(x, bucket_m, bucket_b), 'r')

def discharge(h, t):
    if np.isnan(h).any():
        return np.nan

    elif np.isnan(t).any():
            return np.nan
    else:
        return np.mean(depth_vol(h[0], bucket_m, bucket_b) / t[0])

# slope function is fit to water and bed slope points to extract a single value for flume
def slope(x, S, b):
    return S*x + b

# water density and error
rof = 995 # kg/m^3
drof = 10 # kg/m^3
# gravity
g = 9.8 # m^2/s
# tank width and error
b = 0.0105 # m
db = .001 # error in width measurement


wd = str(pathlib.Path(os.path.dirname(__file__)).parent)
df_density = pd.read_csv(wd + '/1_grain_properties/1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv(wd + '/1_grain_properties/2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)
df_ars = pd.read_csv(wd + '/1_grain_properties/3_var_ARS.csv').set_index('Unnamed: 0').rename_axis(None)
df_meas_settling = pd.read_csv(wd + '/1_grain_properties/4_var_settling_velocities.csv').set_index('Unnamed: 0').rename_axis(None)
df_VES = pd.read_csv(wd + '/1_grain_properties/5_var_VES_settling_velocity.csv').set_index('Unnamed: 0').rename_axis(None)

raw_exp_data = json.load(open(wd + '/2_exp_raw_data/raw_exp_data.txt'))
image_slope_meas = pd.read_csv(wd + '/2_exp_raw_data/image_slope_meas.csv').set_index('Unnamed: 0').rename_axis(None).to_dict()

#%%
database = {}
ns = []
dts = []
list_of_experiments_without_slope_measurements = ['gb_s3_e1', 'gb_s3_e2', 'gb_s3_e3', 'gb_s3_e4', 'gb_s3_e5', 'ng_s3_e2']
for series, dict0 in raw_exp_data.items():
    for name, dict1 in dict0.items():
        gk = dict1['grain_kind']
        D = df_shape['dn'][gk] # volume equivalent sphere diameter (m)
        dD = df_shape['ddn'][gk] # error on VES diam
        ros = df_density['total_grain_ro'][gk] # mean grain density kg/m^3
        dros = df_density['dtotal_grain_ro'][gk] # error on mean grain density

        # normalizing constant for sediment flux
        qs_norm = np.sqrt(g*D*(ros-rof)/rof)*D
        dqs_norm = np.sqrt((g*D**3)/(4*rof*(ros-rof))) * np.sqrt( (3*(ros-rof)*dD/D)**2 + (ros*drof/rof)**2 + dros**2)

        # normalizing constant for shear stress
        tau_norm = (ros - rof)*g*D
        dtau_norm = tau_norm * np.sqrt((dros/ros)**2 + (drof/rof)**2 + (dD/D)**2)

        (h1, h2) = dict1['slope_inds'] # regions of flume to measure slope from
        (m1, m2) = dict1['equil_mass_inds'] # timeframe of total experiment to measure grain mass collected in boxes at end of flume for known period of time
        flume_slope = (dict1['bed_slope']) / 1000  # flume slope (units of grade m/m)
        dflume_slope = 10/1000 # estimated error on measurement of flume slope
        
        L = np.array(dict1['meas_dist_along_flume'])/100 # places where bed height and flow depth was measured along the flume
        h_water_equil = np.array(dict1['water_h_range_equil'])/100 # height of water surface above flume deck as measured during experiment at steady state (m)
        h_bed_equil = np.array(dict1['bed_h_range_equil'])/100 # height of bed surface above flume deck as measured during experiment at steady state (m)

        ns.append(np.array(dict1['bead_mass'][m1:m2]).size)
        dts.append(np.array(dict1['dt_equil'][m1:m2]).mean())
        equil_mass_flux = np.nanmean(np.array(dict1['bead_mass'][m1:m2])/(np.array(dict1['dt_equil'][m1:m2])*1000)) # mass flux in kg/s as measured from end of flume
        dequil_mass_flux = np.nanstd(np.array(dict1['bead_mass'][m1:m2])/(np.array(dict1['dt_equil'][m1:m2])*1000))/np.sqrt(np.array(dict1['bead_mass'][m1:m2]).size)

        qs = equil_mass_flux / (ros*b) # m^2/s # volume flux per unit channel width
        dqs = qs * np.sqrt( (dequil_mass_flux/equil_mass_flux)**2 + (dros/ros)**2 + (db/b)**2 )
        q8 = qs / qs_norm # nondimensional volume flux per channel width
        dq8 = q8 * np.sqrt( (dqs/qs)**2 + (dqs_norm/qs_norm)**2 )

        # if water height was measured, calculate water slope and resulting shear stress (not true for all experiments)
        ### these data were not used
        if series + '_' + name in list_of_experiments_without_slope_measurements:
            # slope meas flag = 1 means that the slopes were measured after the fact from images. This method has the potential to differ slightly from the method used for the majority of the experiments
            slope_meas_flag = 1

            # calculate water/bed slope
            Sw = image_slope_meas[series+'_'+name]['Sw']
            dSw = image_slope_meas[series+'_'+name]['dSw']
            Sb = image_slope_meas[series+'_'+name]['Sb']
            dSb = image_slope_meas[series+'_'+name]['dSb']
            
            # calculate flow depth
            d = image_slope_meas[series+'_'+name]['d'] / 1000
            dd = image_slope_meas[series+'_'+name]['dd'] / 1000
            
            # calculate hydraulic radius in parts to help with error estimation
            R1 = b*d 
            dR1 = R1 * np.sqrt((db/b)**2 + (dd/d)**2)
            R2 = b + 2*d
            dR2 = np.sqrt(db**2 + 4*dd**2)
            R = R1 / R2
            dR = R * np.sqrt((dR1/R1)**2 + (dR2/R2)**2)

            # calculate bed shear stress at equilibrium
            tau_b = np.nan#rof*g*R*Sw # estimated shear stress
            tau_b_bed = np.nan
            tau_b_water = np.nan
            tau_b_corrected = np.nan
            dtau_b = tau_b * np.sqrt((drof/rof)**2 + (dR/R)**2 +(dSw/Sw)**2) # error on shear stress
            tau8 = tau_b / tau_norm # nondimensional shear stress
            dtau8 = tau8 * np.sqrt((dtau_b/tau_b)**2 + (dtau_norm/tau_norm)**2) # error

        ### Only the data below were used
        else:
            # slope meas flag = 0 means that the slopes were measured during the experiment at steady state
            slope_meas_flag = 0
            # calculate equilibrium water slope
            (S_water_equil, m_water_equil), covp_water = op.curve_fit(slope, L[h1:h2], h_water_equil[h1:h2]) # slope relative to flume found by fitting line to water surface height along flume
            dSw_eq = np.sqrt(np.diag(covp_water)[0]) # estimated slope error
            Sw = S_water_equil + flume_slope # actual water slope
            Swb = S_water_equil + flume_slope - 0.01 # estimated bed slope
            dSw = np.sqrt(dSw_eq**2 + dflume_slope**2)
            dSwb = np.sqrt(dSw_eq**2 + dflume_slope**2)

            # calculate equilibrium bed slope
            (S_bed_equil, m_bed_equil), covp_bed = op.curve_fit(slope, L[h1:h2], h_bed_equil[h1:h2]) # slope relative to flume found by fitting line to bed surface height along flume
            dSb_eq = np.sqrt(np.diag(covp_bed)[0])
            Sb = S_bed_equil + flume_slope # actual bed slope
            dSb = np.sqrt(dSb_eq**2 + dflume_slope**2)

            # calculate equilibrium water depth and its mean and standard deviation
            d_equils = h_water_equil[h1:h2] - h_bed_equil[h1:h2]
            d = d_equils.mean()
            dd = d_equils.std()/np.sqrt(d_equils.size)

            # calculate hydraulic radius in parts to help with error estimation
            R1 = b*d 
            dR1 = R1 * np.sqrt((db/b)**2 + (dd/d)**2)
            R2 = b + 2*d
            dR2 = np.sqrt(db**2 + 4*dd**2)
            R = R1 / R2
            dR = R * np.sqrt((dR1/R1)**2 + (dR2/R2)**2)

            Q = discharge(dict1['discharge_depth'], dict1['discharge_time']) / 1000
            So = Sw
            dSo = dSw
            hn = d
            dhn = dd
            W = b
            dW = db
            nu = 1e-6
            # hydraulic radius
            Rh = R
            dRh = dR
            # flow velocity
            V = Q/W/hn
            dV = V * np.sqrt((dW/W)**2 + (dhn/hn)**2)
            # Reynolds number
            Re = 4*V*Rh/nu
            dRe = Re * np.sqrt((dV/V)**2 + (dRh/Rh)**2)
            # shear velocity
            u8 = np.sqrt(9.81 * Rh * So)
            du8 = (u8/2) * np.sqrt((dRh/Rh)**2 + (dSo/So)**2)
            # darcy weisbach factor
            f = 8*u8**2/V**2
            d_f = 2*f * np.sqrt((du8/u8)**2 + (dV/V)**2)
            # Wall friction factor
            fw = 0.0026 * (np.log10(Re/f))**2 - 0.0428*(np.log10(Re/f)) + 0.1884
            dfw = (2*0.0026*np.log(Re/f) - 0.0428) * np.sqrt((dRe/Re)**2 + (d_f/f)**2)
            # bed friction factor
            fb = f + 2*(hn/W)*(f-fw)
            dfb = fb*np.sqrt((dhn/hn)**2 + (dW/W)**2 + (dfw/fw)**2 + (d_f/f)**2)
            # bed hydraulic radius
            Rb = fb/f * Rh
            dRb = Rb * np.sqrt((dRh/Rh)**2 + (d_f/f)**2 + (dfb/fb)**2)
            tau_b_corrected = rof*g*Rb*So # estimated shear stress
            dtau_b_corrected = tau_b_corrected * np.sqrt((dRb/Rb)**2 + (dSo/So)**2)

            # calculate bed shear stress at equilibrium
            tau_b = rof*g*R*Swb # estimated shear stress
            tau_b_water = rof*g*R*Sw # estimated shear stress
            tau_b_bed = rof*g*R*Sb # estimated shear stress
            dtau_b = tau_b * np.sqrt((drof/rof)**2 + (dR/R)**2 +(dSw/Sw)**2) # error on shear stress
            tau8 = tau_b / tau_norm # nondimensional shear stress
            dtau8 = tau8 * np.sqrt((dtau_b/tau_b)**2 + (dtau_norm/tau_norm)**2) # error
            tau8_corrected = tau_b_corrected / tau_norm
            dtau8_corrected = tau8_corrected * np.sqrt((dtau_b_corrected/tau_b_corrected)**2 + (dtau_norm/tau_norm)**2) # error
            # print(dSo/So, dhn/hn, dW/W, dRh/Rh, dd/d, dV/V, dRe/Re, du8/u8, d_f/f, dfw/fw, dfb/fb, dRb/Rb)


        # save all parameters for each experiment as a dict in the dict of dicts 'database'
        database[series+'_'+name] = {
            'names': series+'_'+name,
            'series': series[-1],
            'exp': name[-1],
            'grain_kind': gk,
            'flume_slope': dict1['bed_slope'],
            'VES_grain_diameter': D,
            'dVES_grain_diameter': dD,
            'discharge': discharge(dict1['discharge_depth'], dict1['discharge_time']),
            'grain_vol': df_density['grain_vol'][gk],
            'dgrain_vol': df_density['dgrain_vol'][gk],
            'grain_mass': df_density['grain_mass'][gk],
            'dgrain_mass': df_density['dgrain_mass'][gk],
            'VES_settling_velocity': df_VES['ws_sphere'][gk],
            'dVES_settling_velocity': df_VES['dws_sphere'][gk],
            'measured_settling_velocity': df_meas_settling['vel_mean'][gk],
            'dmeasured_settling_velocity': df_meas_settling['dvel_mean'][gk],
            'angle_of_repose': df_ars['mean_angle'][gk],
            'dangle_of_repose': df_ars['std_angle'][gk],
            'grain_shape_A': df_shape['A'][gk],
            'dgrain_shape_A': df_shape['dA'][gk],
            'grain_shape_B': df_shape['B'][gk],
            'dgrain_shape_B': df_shape['dB'][gk],
            'grain_shape_C': df_shape['C'][gk],
            'dgrain_shape_C': df_shape['dC'][gk],
            'grain_shape_corey_shape_factor': df_shape['CSF'][gk],
            'dgrain_shape_corey_shape_factor': df_shape['dCSF'][gk],
            'grain_density': ros,
            'dgrain_density': dros,
            'fluid_density': rof,
            'dfluid_density': drof,
            'mass_flux': equil_mass_flux,
            'dmass_flux': dequil_mass_flux,
            'volume_flux': qs,
            'dvolume_flux': dqs,
            'q8': q8,
            'dq8': dq8,
            'taub': tau_b,
            'taub_bed': tau_b_bed,
            'taub_water': tau_b_water,
            'taub_corrected': tau_b_corrected,
            'dtaub_corrected': dtau_b_corrected,
            'dtaub': dtau_b,
            'tau8': tau8,
            'dtau8': dtau8,
            'tau8_corrected': tau8_corrected,
            'dtau8_corrected': dtau8_corrected,
            'hyd_rad': R,
            'dhyd_rad': dR,
            'flow_depth': d,
            'dflow_depth': dd,
            'water_slope': Sw,
            'dwater_slope': dSw,
            'bed_slope': Sb,
            'dbed_slope': dSb,
            'slope_meas_flag': slope_meas_flag
        }

df = pd.DataFrame(database).T
df.to_csv('exp_data.txt')

# %%
